146a91
@@ -15,7 +15,9 @@
 package org.apache.hadoop.hive.llap.tezplugins;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Comparator;
 import java.util.EnumSet;
 import java.util.HashMap;
@@ -781,27 +783,26 @@
private SelectHostResult selectHost(TaskInfo request) {
         }
       }
       /* fall through - miss in locality (random scheduling) or no locality-requested */
-      Entry<String, NodeInfo>[] all = instanceToNodeMap.entrySet().toArray(new Entry[0]);
-      // Check again
+      Collection<ServiceInstance> instances = activeInstances.getAll();
+      ArrayList<NodeInfo> all = new ArrayList<>(instances.size());
+      for (ServiceInstance inst : instances) {
+        NodeInfo nodeInfo = instanceToNodeMap.get(inst.getWorkerIdentity());
+        if (nodeInfo != null && nodeInfo.canAcceptTask()) {
+          all.add(nodeInfo);
+        }
+      }
       if (LOG.isDebugEnabled()) {
         LOG.debug("Attempting random allocation for task={}", request.task);
       }
-      if (all.length > 0) {
-        int n = random.nextInt(all.length);
-        // start at random offset and iterate whole list
-        for (int i = 0; i < all.length; i++) {
-          Entry<String, NodeInfo> inst = all[(i + n) % all.length];
-          if (inst.getValue().canAcceptTask()) {
-            LOG.info(
-                "Assigning " + nodeToString(inst.getValue().getServiceInstance(), inst.getValue()) +
-                    " when looking for any host, from #hosts=" + all.length + ", requestedHosts=" +
-                    ((requestedHosts == null || requestedHosts.length == 0) ? "null" :
-                        Arrays.toString(requestedHosts)));
-            return new SelectHostResult(inst.getValue().getServiceInstance(), inst.getValue());
-          }
-        }
+      if (all.isEmpty()) {
+        return SELECT_HOST_RESULT_DELAYED_RESOURCES;
       }
-      return SELECT_HOST_RESULT_DELAYED_RESOURCES;
+      NodeInfo randomNode = all.get(random.nextInt(all.size()));
+      LOG.info("Assigning " + nodeToString(randomNode.getServiceInstance(), randomNode)
+          + " when looking for any host, from #hosts=" + all.size() + ", requestedHosts="
+          + ((requestedHosts == null || requestedHosts.length == 0)
+              ? "null" : Arrays.toString(requestedHosts)));
+      return new SelectHostResult(randomNode.getServiceInstance(), randomNode);
     } finally {
       readLock.unlock();
     }
